import os
from modelscope import snapshot_download
from tqdm import tqdm


def download_qwen3_4b():
    """从魔搭社区下载Qwen3-4B模型"""

    # 魔搭社区上的Qwen3-4B模型地址
    model_name = "Qwen/Qwen3-4B"
    model_dir = "./qwen3-4b-model"

    # 创建目录
    os.makedirs(model_dir, exist_ok=True)

    print(f"开始从魔搭社区下载模型: {model_name}")
    print("下载可能需要一些时间，请耐心等待...")

    try:
        # 使用modelscope下载模型
        # snapshot_download会自动处理缓存和断点续传
        model_path = snapshot_download(
            model_name,
            cache_dir=model_dir,
            revision='master'  # 默认使用主分支
        )

        print(f"模型下载完成！保存路径: {model_path}")
        return model_path

    except Exception as e:
        print(f"下载过程中出现错误: {e}")
        return None


def download_qwen3_4b_instruct():
    """下载Qwen3-4B-Instruct版本（如果需要指令调优版本）"""

    model_name = "Qwen/Qwen3-4B-Instruct"
    model_dir = "./qwen3-4b-instruct-model"

    os.makedirs(model_dir, exist_ok=True)

    print(f"开始下载指令调优版本: {model_name}")

    try:
        model_path = snapshot_download(
            model_name,
            cache_dir=model_dir,
            revision='master'
        )

        print(f"指令调优版本下载完成！保存路径: {model_path}")
        return model_path

    except Exception as e:
        print(f"下载过程中出现错误: {e}")
        return None


def check_download_progress():
    """检查下载进度（基础版本）"""
    model_dir = "./qwen3-4b-model"

    if os.path.exists(model_dir):
        # 计算目录大小
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(model_dir):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                total_size += os.path.getsize(fp)

        # 转换为GB
        size_gb = total_size / (1024 ** 3)
        print(f"当前已下载: {size_gb:.2f} GB")

        # Qwen3-4B大约7-8GB，可以根据这个估算进度
        if size_gb > 6:
            print("下载接近完成！")
        elif size_gb > 3:
            print("下载进行中...")
        else:
            print("下载刚开始...")

    return model_dir


if __name__ == "__main__":
    # 首先确保安装了modelscope
    try:
        import modelscope
    except ImportError:
        print("请先安装modelscope: pip install modelscope")
        exit(1)

    # 下载基础版本
    model_path = download_qwen3_4b()

    if model_path:
        print("\n" + "=" * 50)
        print("下载成功！您可以使用以下方式加载模型：")
        print("=" * 50)
        print("\n使用transformers加载：")
        print(f"from transformers import AutoModel, AutoTokenizer")
        print(f"model = AutoModel.from_pretrained('{model_path}')")
        print(f"tokenizer = AutoTokenizer.from_pretrained('{model_path}')\n")

        print("或者使用modelscope加载：")
        print(f"from modelscope import AutoModel, AutoTokenizer")
        print(f"model = AutoModel.from_pretrained('{model_path}')")
        print(f"tokenizer = AutoTokenizer.from_pretrained('{model_path}')")

    # 可选：下载指令调优版本
    download_instruct = input("\n是否同时下载指令调优版本？(y/n): ").lower().strip()
    if download_instruct == 'y':
        download_qwen3_4b_instruct()